       
       
       
        seg_logits = resize(
            input=seg_logits,
            size=gt_semantic_seg.shape[2:],
            mode='bilinear',
            align_corners=False)
        
        
        # unorm = UnNormalize(mean=(123.675, 116.28, 103.53), std=(58.395, 57.12, 57.375))
        mean = [123.675, 116.28, 103.53]
        std = [58.395, 57.12, 57.375]
        unorm = UnNormalize(mean=mean, std=std)

        seg_logits_np = seg_logits.detach().cpu().numpy()
        gt_semantic_seg_np = gt_semantic_seg.detach().cpu().numpy()

        # Assuming seg_logits has shape (N, C, H, W) and gt_semantic_seg has shape (N, 1, H, W)
        # We'll take the argmax along the channel dimension for seg_logits for visualization
        img_save_path = '/data1/users/zhengzhiyu/mtp_workplace/documents/OBBDetection/work_dirs/visualization/img/'
        mask_save_path = '/data1/users/zhengzhiyu/mtp_workplace/documents/OBBDetection/work_dirs/visualization/mask/'
        pred_save_path = '/data1/users/zhengzhiyu/mtp_workplace/documents/OBBDetection/work_dirs/visualization/pred/'
        seg_preds_np = np.argmax(seg_logits_np, axis=1)

        for i in range(seg_preds_np.shape[0]):
            image = unorm(img[i].clone()).cpu().numpy()
            image = image.transpose(1, 2, 0)
            image = np.clip(image, 0, 255).astype(np.uint8)

            seg_pred_img_rgb = colors[seg_preds_np[i]].astype(np.uint8)
            gt_seg_img_rgb = colors[gt_semantic_seg_np[i, 0]].astype(np.uint8)

            seg_pred_img = Image.fromarray(seg_pred_img_rgb)
            gt_seg_img = Image.fromarray(gt_seg_img_rgb)
            image = Image.fromarray(image)
            file_name = img_metas[0]['ori_filename'][:-4]
            seg_pred_path = os.path.join(img_save_path, f'seg_pred_{file_name}.png')
            gt_seg_path = os.path.join(mask_save_path, f'gt_seg_{file_name}.png')
            image_path = os.path.join(pred_save_path, f'image_{file_name}.png')

            seg_pred_img.save(seg_pred_path)
            gt_seg_img.save(gt_seg_path)
            image.save(image_path)
            # print(f'Saved seg_pred to {seg_pred_path} and gt_seg to {gt_seg_path}')